import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models.resnet import resnet101

class ResNet101_FPN(nn.Module):
    def __init__(self):
        super().__init__()

        # 加载ResNet101主干（关键修改点）
        resnet = resnet101(pretrained=True)

        # 修改第一层卷积（适配单通道输入）
        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        with torch.no_grad():  # 初始化权重（使用RGB三通道的均值）
            self.conv1.weight.copy_(resnet.conv1.weight.mean(dim=1, keepdim=True))

        # 继承其他层结构
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool

        # ResNet101的阶段层（注意：layer结构名称与ResNet50相同）
        self.layer1 = resnet.layer1  # 输出256通道
        self.layer2 = resnet.layer2  # 输出512通道
        self.layer3 = resnet.layer3  # 输出1024通道

        # ----------------------------
        # FPN结构（与ResNet50版本完全兼容）
        # ----------------------------
        self.lat3 = nn.Conv2d(1024, 256, 1)  # 输入通道与ResNet101的layer3输出一致
        self.lat2 = nn.Conv2d(512, 256, 1)
        self.lat1 = nn.Conv2d(256, 128, 1)

        self.smooth3 = nn.Conv2d(256, 256, 3, padding=1)
        self.smooth2 = nn.Conv2d(128, 128, 1)

        self.smooth1 = nn.Conv2d(256, 128, 1)
    def forward(self, x):
        # 输入尺寸 [B, 1, 640, 480]
        x = self.conv1(x)  # [B, 64, 320, 240]
        x = self.bn1(x)
        x = self.relu(x)
        c1 = self.maxpool(x)  # [B, 64, 160, 120]

        c2 = self.layer1(c1)  # [B, 256, 160, 120]
        c3 = self.layer2(c2)  # [B, 512, 80, 60]
        c4 = self.layer3(c3)  # [B, 1024, 40, 30]

        # FPN融合流程（与ResNet50相同）
        p4 = self.lat3(c4)  # [B, 256, 40, 30]
        # p3 = self._upsample_add(p4, self.lat2(c3))  # [B, 256, 80, 60]
        p4 = F.interpolate(p4, scale_factor=2, mode='bilinear', align_corners=True)
        p3 = p4 + self.lat2(c3)

        p3 = self.smooth3(p3)  # p8输出

        # p2 = self._upsample_add(p3, self.lat1(c2))  # [B, 128, 160, 120]
        p3 = F.interpolate(p3, scale_factor=2, mode='bilinear', align_corners=True)
        p2 = self.smooth1(p3) + self.lat1(c2)

        p2 = self.smooth2(p2)
        p2 = F.interpolate(p2, scale_factor=2, mode='bilinear')  # [B, 128, 320, 240]

        return [p3, p2]  # [(256,80,60), (128,320,240)]

model = ResNet101_FPN()
input_tensor = torch.randn(1, 1, 640, 480)
p8, p2 = model(input_tensor)

print(p8.shape)  # torch.Size([1, 256, 80, 60])
print(p2.shape)  # torch.Size([1, 128, 320, 240])